{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "\n",
    "This notebook documents test suites in Checklist. If you are not already familiar with creating tests in Checklist, consider reading the MFT Examples notebook.\n",
    "\n",
    "## Setup\n",
    "First, let's import the libraries and load the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
    "import pandas as pd\n",
    "import random\n",
    "import json\n",
    "import checklist\n",
    "from checklist.editor import Editor\n",
    "from checklist.expect import Expect\n",
    "from checklist.pred_wrapper import PredictorWrapper\n",
    "from checklist.test_types import MFT\n",
    "from checklist.test_suite import TestSuite\n",
    "from torch.nn import functional as F\n",
    "from typing import List\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fe4e8058110>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize random seed\n",
    "# Remove this code to experiment with random samples\n",
    "random.seed(123)\n",
    "torch.manual_seed(456)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Model loaded'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load pretrained model tokenizer\n",
    "tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "# Load pretrained model (weights)\n",
    "model = GPT2LMHeadModel.from_pretrained(\"gpt2\", pad_token_id=tokenizer.eos_token_id)\n",
    "device = 'cuda'\n",
    "model.eval()\n",
    "model.to(device)\n",
    "\"Model loaded\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a Test Suite\n",
    "\n",
    "Checklist can run multiple tests in a test suite. Tests can be grouped by capability and results can be explored in a visual table.\n",
    "\n",
    "We will create a test suite called 'Same Token Prediction' with 3 MFTs. Each MFT will test if the token substituted into the prompt template also appears in the generated text.\n",
    "\n",
    "For example, if we prompt the model with \"The **dog** is running in the zoo\" and the model responds with \"The **dog** looks very happy\", then it passes the test because the same animal appears in the model's response.\n",
    "\n",
    "## Creating the MFTs\n",
    "### MFT 1: Same animal appears in response\n",
    "This MFT uses an `{animal}` placeholder in the template. The expectation function checks that the same animal appears in the prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['The dog is running in the zoo',\n",
       " 'The cat is running in the zoo',\n",
       " 'The giraffe is running in the zoo',\n",
       " 'The aardvark is running in the zoo']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "editor = Editor()\n",
    "animal_prompts = editor.template(\"The {animal} is running in the zoo\", animal=[\"dog\", \"cat\", \"giraffe\", \"aardvark\"], meta=True)\n",
    "animal_prompts.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def contains_same_animal(x, pred, conf, label=None, meta=None):\n",
    "    return meta['animal'] in pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "same_animal_expect_fn = Expect.single(contains_same_animal)\n",
    "same_animal_test = MFT(**animal_prompts, name='Same animal in response', description='The response contains the same animal mentioned in the prompt.', expect=same_animal_expect_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MFT 2: Same country appears in response\n",
    "This MFT uses a `{country}` placeholder in the template. The expectation function checks that the same country appears in the prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MunchWithAdd({'meta': [{'country': 'Cambodia'}, {'country': 'Bangladesh'}, {'country': 'Spain'}, {'country': 'Solomon Islands'}, {'country': 'Eritrea'}, {'country': 'Togo'}, {'country': 'Mexico'}, {'country': 'Sudan'}, {'country': 'Saint Kitts and Nevis'}, {'country': 'India'}], 'data': ['Earlier today, scientists from Cambodia discovered', 'Earlier today, scientists from Bangladesh discovered', 'Earlier today, scientists from Spain discovered', 'Earlier today, scientists from Solomon Islands discovered', 'Earlier today, scientists from Eritrea discovered', 'Earlier today, scientists from Togo discovered', 'Earlier today, scientists from Mexico discovered', 'Earlier today, scientists from Sudan discovered', 'Earlier today, scientists from Saint Kitts and Nevis discovered', 'Earlier today, scientists from India discovered']})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "country_prompts = editor.template(\"Earlier today, scientists from {country} discovered\", meta=True, nsamples=10)\n",
    "country_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def contains_same_country(x, pred, conf, label=None, meta=None):\n",
    "    return meta['country'] in pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "same_country_expect_fn = Expect.single(contains_same_country)\n",
    "same_country_test = MFT(**country_prompts, name='Same country in response', description='The response contains the same country mentioned in the prompt.', expect=same_country_expect_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MFT 3: Same person appears in response\n",
    "This MFT uses a `{first_name}` placeholder in the template. The expectation function checks that the same first name appears in the prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MunchWithAdd({'meta': [{'first_name': 'Carolyn'}, {'first_name': 'John'}, {'first_name': 'George'}, {'first_name': 'Lawrence'}, {'first_name': 'Linda'}, {'first_name': 'Steve'}, {'first_name': 'Judy'}, {'first_name': 'Alfred'}, {'first_name': 'Patricia'}, {'first_name': 'Jill'}], 'data': ['Carolyn is my neighbor.', 'John is my neighbor.', 'George is my neighbor.', 'Lawrence is my neighbor.', 'Linda is my neighbor.', 'Steve is my neighbor.', 'Judy is my neighbor.', 'Alfred is my neighbor.', 'Patricia is my neighbor.', 'Jill is my neighbor.']})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "person_prompts = editor.template(\"{first_name} is my neighbor.\", meta=True, nsamples=10)\n",
    "person_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def contains_same_person(x, pred, conf, label=None, meta=None):\n",
    "    return meta['first_name'] in pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "same_person_expect_fn = Expect.single(contains_same_person)\n",
    "same_person_test = MFT(**person_prompts, name='Same person in response', description='The response contains the same person\\'s first name mentioned in the prompt.', expect=same_person_expect_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding the tests to the suite\n",
    "The `TestSuite()` constructor creates an empty test suite. Tests can be added one by one using `suite.add(test)`. The optional `capability` parameter can be used to label and group tests that test similar capabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite = TestSuite()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite.add(same_animal_test, capability=\"Same Token Prediction\")\n",
    "suite.add(same_country_test, capability=\"Same Token Prediction\")\n",
    "suite.add(same_person_test, capability=\"Same Token Prediction\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the predictions\n",
    "Now we define the function that Checklist will use to generate predictions from the model. The predictions need to be returned in the form `([predictions], [scores])`, so we will wrap the `generate_sentences()` function with `PredictorWrapper.wrap_predict()` to automatically create a tuple `([predictions], [1, 1, ...])`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_sentence(prompt: str) -> str:\n",
    "    token_tensor = tokenizer.encode(prompt, return_tensors='pt').to(device) # return_tensors = \"pt\" returns a PyTorch tensor\n",
    "    out = model.generate(\n",
    "        token_tensor,\n",
    "        do_sample=True,\n",
    "        min_length=10,\n",
    "        max_length=50,\n",
    "        num_beams=1,\n",
    "        temperature=1.0,\n",
    "        no_repeat_ngram_size=2,\n",
    "        early_stopping=False,\n",
    "        output_scores=True,\n",
    "        return_dict_in_generate=True)\n",
    "    text = tokenizer.decode(out.sequences[0], skip_special_tokens=True)\n",
    "    return text[len(prompt):]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_sentences(prompts: List[str]) -> List[str]:\n",
    "    sentences = []\n",
    "    for prompt in prompts:\n",
    "        sentences.append(generate_sentence(prompt))\n",
    "    return sentences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([' Now, let us begin to talk…\\n\\nI heard you guys were coming.\\n: I thought you were from there, right?\\n (Slight sigh, looks back at Vixen) Yeah,',\n",
       "  ' Goodbye, goodbye, dear.\" But when I heard that I felt almost lost.\\n\\nMy wife gave me the same answer of \"Thank you.\"\\n.'],\n",
       " array([1., 1.]))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wrapped_generator = PredictorWrapper.wrap_predict(generate_sentences)\n",
    "wrapped_generator([\"Hello, nice to meet you.\", \"Goodbye, see you later.\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the suite\n",
    "We can now run the suite and view the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Same animal in response\n",
      "Predicting 4 examples\n",
      "Running Same country in response\n",
      "Predicting 10 examples\n",
      "Running Same person in response\n",
      "Predicting 10 examples\n"
     ]
    }
   ],
   "source": [
    "suite.run(wrapped_generator, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_example(x, pred, conf, label=None, meta=None): \n",
    "    return 'Prompt:      %s\\nCompletion:      %s' % (x, pred) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Same Token Prediction\n",
      "\n",
      "Same animal in response\n",
      "Test cases:      4\n",
      "Fails (rate):    4 (100.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      The dog is running in the zoo\n",
      "Completion:      , she says, and a friend of her walks in on them, shouting things like, \"I was scared, can you please sit down\"?\n",
      "\n",
      "He's also apparently had a seizure and an eye scratch,\n",
      "----\n",
      "Prompt:      The cat is running in the zoo\n",
      "Completion:      , but I'll only bring back a one year old. There's one guy at the park that is crazy. He was doing some nice things but it's definitely going to end up being one of those stories that\n",
      "----\n",
      "Prompt:      The aardvark is running in the zoo\n",
      "Completion:      . I got to see it in person.\n",
      "\n",
      "Lately, my daughter, who was visiting from college now, said she's been fascinated with the penguin-shaped shell on the side of\n",
      "----\n",
      "\n",
      "\n",
      "Same country in response\n",
      "Test cases:      10\n",
      "Fails (rate):    9 (90.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Earlier today, scientists from Spain discovered\n",
      "Completion:       that hydrogen sulfide makes a \"crystalline transition\" into a metal. This is the first time that a chemical reaction can be produced when both the organic matter and hydrocarbons are oxidized to form hydrogen\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Mexico discovered\n",
      "Completion:       its first \"sterile\" fish, and it isn't even the first time that some species have made it back home from the Mediterranean.\n",
      "\n",
      "\"This study finds the oldest possible species of marine mammal came back\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Eritrea discovered\n",
      "Completion:       a new species of squid called Staphylococcus aureus and are planning to use genetically modified genes to identify its killer.\n",
      "\n",
      "A study published in the journal BMC Genetics says that a unique species\n",
      "----\n",
      "\n",
      "\n",
      "Same person in response\n",
      "Test cases:      10\n",
      "Fails (rate):    10 (100.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Lawrence is my neighbor.\n",
      "Completion:       I'm just glad that you haven't got a message of disapproval.\"\n",
      "\n",
      "He added: \"This whole point of this project, this is about sharing it with other Canadians who don't know what it means and aren\n",
      "----\n",
      "Prompt:      John is my neighbor.\n",
      "Completion:       I have no idea how he got here [to help] these poor folks up their own tracks, to take on this sort of monstrous company of wolves.\"\n",
      "\n",
      "Wes told The Christian Science Monitor that the idea has inspired\n",
      "----\n",
      "Prompt:      Alfred is my neighbor.\n",
      "Completion:       His name has become synonymous with a man of many names. He is so strong and so fast that if he doesn't fight for his wife, I never know of a time when that would be possible.\n",
      "\n",
      "\n",
      "----\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "suite.summary(format_example_fn = format_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please wait as we prepare the table data...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f01ba2c6f59467b93bd81d310302fa8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "SuiteSummarizer(stats={'npassed': 0, 'nfailed': 0, 'nfiltered': 0}, test_infos=[{'name': 'Same animal in respo…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "suite.visual_summary_table()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using files with test suites\n",
    "\n",
    "Some models cannot be run directly on the same machine that is running the Checklist test suite. For instance, a model might need to run in a specially configured lab environment. In this case, Checklist does not have to receive the predictions from the model directly. The predictions can be saved to a file, then the test suite can check the predictions from the file.\n",
    "\n",
    "## Exporting a test suite to a file\n",
    "First, let's create a file that contains all the prompts that we will send to the model. TestSuite's `to_raw_file()` function exports a test suite to a file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite.to_raw_file(\"suite.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The dog is running in the zoo\r\n",
      "The cat is running in the zoo\r\n",
      "The giraffe is running in the zoo\r\n",
      "The aardvark is running in the zoo\r\n",
      "Earlier today, scientists from Cambodia discovered\r\n",
      "Earlier today, scientists from Bangladesh discovered\r\n",
      "Earlier today, scientists from Spain discovered\r\n",
      "Earlier today, scientists from Solomon Islands discovered\r\n",
      "Earlier today, scientists from Eritrea discovered\r\n",
      "Earlier today, scientists from Togo discovered\r\n",
      "Earlier today, scientists from Mexico discovered\r\n",
      "Earlier today, scientists from Sudan discovered\r\n",
      "Earlier today, scientists from Saint Kitts and Nevis discovered\r\n",
      "Earlier today, scientists from India discovered\r\n",
      "Carolyn is my neighbor.\r\n",
      "John is my neighbor.\r\n",
      "George is my neighbor.\r\n",
      "Lawrence is my neighbor.\r\n",
      "Linda is my neighbor.\r\n",
      "Steve is my neighbor.\r\n",
      "Judy is my neighbor.\r\n",
      "Alfred is my neighbor.\r\n",
      "Patricia is my neighbor.\r\n",
      "Jill is my neighbor."
     ]
    }
   ],
   "source": [
    "cat \"suite.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the predictions\n",
    "Next, we need to generate predictions and save them to another file named `suite_predictions.txt`. By default, Checklist expects to receive 1 prediction per line, so we will replace any newline characters with spaces in each prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('suite.txt', 'r') as input_file:\n",
    "    with open('suite_predictions.txt', 'w') as output_file:\n",
    "        for line in input_file:\n",
    "            prediction = generate_sentence(line)\n",
    "            prediction = prediction.replace('\\n', ' ')\n",
    "            print(prediction, file=output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Police said the dog and the driver are to be charged with aggravated dangerous driving, after photos surfaced of the woman struggling and being restrained. , which was brought to police in August, was found unconscious while\r\n",
      " If you're new to dogfighting, you might be familiar with the situation that happened last Tuesday — a big chase at a wildlife conservancy in Ohio with nearly 100 dogs in different colors.  (The\r\n",
      " A giraffey tiger is seen inside a zoo in China. The zoo, which is managed by the National Geographic Society, has been hit by one of the biggest conservation disasters in its history, when\r\n",
      " It's one of our favorite things to do. We have all these people running around the park, and they have a special place in our hearts. I love it. , the great\r\n",
      " two distinct new molecules that can react with each other to form new genes. These molecules make use of two proteins — the protein X and the non-protein X2 called the X chromosome — to coordinate their\r\n",
      " more than 4,000 new species of amphibians  'that are being observed globally and beyond,' said scientists at the Bangladesh Biosphere Institute.  of more than 14,300 species found in the world\r\n",
      " the first evidence of a human hybridization of two species—three of which have human characteristics, such as white skin and blond hair—at the beginning of the Pleistocene epoch.  'This has\r\n",
      " the first new fossil discovered for the first time at Pluto's northern limb. The discovery, made in 2010, demonstrates that the surface of the moon may harbour . , a new planet-sized\r\n",
      " they believe the human brain has been altered by an act of God, with the result of a brain injury.  (AFP/GETTY) ... The Human Cell Project and the University\r\n",
      " a new type of protein to keep human cells from dying inside. It is now known to be in the blood of the most .    \"The researchers discovered an important new mechanism by which\r\n",
      " a major new finding that would make them the first to show —as the researchers said a major breakthrough .  The research showed that the process  — also known as  'finite entropy\r\n",
      " A team of Egyptian scientists were able to test out new technologies for making small, lighter, more efficient lasers which would work on thin materials. ... a very interesting technique that would change the way\r\n",
      " some sort of planet-sized planet to be in the Milky Way , that is at its absolute end, approximately 11,999 light years away. The tiny star, which is\r\n",
      " an ancient, three-dimensional human brain : In a team led by researchers at the National Institute of Mental Health (IMH) in Ramakrishnan in Maharashtra, the team had detected the unique\r\n",
      " What do you think of her? Is there a point where having a wife is the way to go for you?  10 Best Ways to Spend a Week (and a Day) With Carolyn Breen /\r\n",
      " This morning I decided to leave my house to get some rest. I was told that I'd missed her, especially the afternoon that her husband drove to the house a little earlier. When I got up I found that\r\n",
      " A great friend and a wonderful person, Mr. George has helped to create many incredible new items. From the large box of ice cubes to his amazing handmade bag of all things, he is like a kid being raised\r\n",
      " \"I went into the house because I was so frustrated, like, 'What are these things happening to my house?' And then I realized, at that point I'm not that upset I didn't ask for\r\n",
      " What if there are only three people in the room!? -Is it alright?! The man was very quiet. He did not tell the people there. But at the same time, the situation made the conversation\r\n",
      " Luna has come here and will stay. I want you to let it fall on your head that she is the only one who is going to come down here at all. Just let us watch the other kids,\r\n",
      " He had a white shirt on and black pants. A black glove on his hip. Two long black belts. He had two gray hair on (yes, black hair) He was wearing a dark green sweater,\r\n",
      " I can't even help myself from looking deeply at some of his pictures and trying to guess what he's really like or what the hell he does for a living. His father always seemed to be just about\r\n",
      " I have to say that she's always been pretty nice to me. She's a very helpful person and I really like her being around other people's children. I love the fact that I am a good mother\r\n",
      " We talked about that yesterday, and he said it just feels awesome and like, I have this awesome time with my kids. I think the kids have it right.  What happened to you when the NFL started playing\r\n"
     ]
    }
   ],
   "source": [
    "cat \"suite_predictions.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the suite from the file\n",
    "Now all we need to do is use `run_from_file()` to check if the predictions pass the tests. Since we don't care about confidence scores in these tests, we will tell checklist to ignore them by specifying `file_format=\"pred_only\"`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite.run_from_file(\"suite_predictions.txt\", file_format=\"pred_only\", overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Same Token Prediction\n",
      "\n",
      "Same animal in response\n",
      "Test cases:      4\n",
      "Fails (rate):    2 (50.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      The aardvark is running in the zoo\n",
      "Completion:       It's one of our favorite things to do. We have all these people running around the park, and they have a special place in our hearts. I love it. , the great\n",
      "----\n",
      "Prompt:      The cat is running in the zoo\n",
      "Completion:       If you're new to dogfighting, you might be familiar with the situation that happened last Tuesday — a big chase at a wildlife conservancy in Ohio with nearly 100 dogs in different colors.  (The\n",
      "----\n",
      "\n",
      "\n",
      "Same country in response\n",
      "Test cases:      10\n",
      "Fails (rate):    9 (90.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Earlier today, scientists from Cambodia discovered\n",
      "Completion:       two distinct new molecules that can react with each other to form new genes. These molecules make use of two proteins — the protein X and the non-protein X2 called the X chromosome — to coordinate their\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Eritrea discovered\n",
      "Completion:       they believe the human brain has been altered by an act of God, with the result of a brain injury.  (AFP/GETTY) ... The Human Cell Project and the University\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Saint Kitts and Nevis discovered\n",
      "Completion:       some sort of planet-sized planet to be in the Milky Way , that is at its absolute end, approximately 11,999 light years away. The tiny star, which is\n",
      "----\n",
      "\n",
      "\n",
      "Same person in response\n",
      "Test cases:      10\n",
      "Fails (rate):    8 (80.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Linda is my neighbor.\n",
      "Completion:       What if there are only three people in the room!? -Is it alright?! The man was very quiet. He did not tell the people there. But at the same time, the situation made the conversation\n",
      "----\n",
      "Prompt:      Alfred is my neighbor.\n",
      "Completion:       I can't even help myself from looking deeply at some of his pictures and trying to guess what he's really like or what the hell he does for a living. His father always seemed to be just about\n",
      "----\n",
      "Prompt:      Steve is my neighbor.\n",
      "Completion:       Luna has come here and will stay. I want you to let it fall on your head that she is the only one who is going to come down here at all. Just let us watch the other kids,\n",
      "----\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "suite.summary(format_example_fn = format_example)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced file example: JSON format\n",
    "Sometimes, we might need some more control over how Checklist imports and exports the examples and predictions from the files. For example, we might want to use JSON formatted data so that we can add more information to each example. In this tutorial, we will add an `id` to each example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exporting to JSON file with to_raw_file()\n",
    "The `format_fn` parameter allows us to control how each example in the suite is printed to the file. We can use `format_fn` to print the examples in a JSON format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def suite_to_json_file(suite, filename):\n",
    "    class Counter:\n",
    "        def __init__(self):\n",
    "            self.count = 0\n",
    "        def get_count(self):\n",
    "            self.count += 1\n",
    "            return self.count\n",
    "    \n",
    "    counter = Counter()\n",
    "    total_tests = 0\n",
    "    for t in suite.tests.values():\n",
    "        total_tests += len(t.data)\n",
    "        \n",
    "    def json_format_fn(x):\n",
    "        example_id = counter.get_count()\n",
    "        json_str = \"\"\n",
    "        if example_id == 1:\n",
    "            json_str = '{\"examples\": ['\n",
    "        json_str += json.dumps({'content': x, 'id': example_id}) + \",\"\n",
    "        if example_id == total_tests:\n",
    "            # remove trailing comma\n",
    "            json_str = json_str[:len(json_str)-1]\n",
    "            json_str += \"]}\"\n",
    "        return json_str\n",
    "    \n",
    "    suite.to_raw_file(filename, format_fn = json_format_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite_to_json_file(suite, 'same_token_suite.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"examples\": [{\"content\": \"The dog is running in the zoo\", \"id\": 1},\r\n",
      "{\"content\": \"The cat is running in the zoo\", \"id\": 2},\r\n",
      "{\"content\": \"The giraffe is running in the zoo\", \"id\": 3},\r\n",
      "{\"content\": \"The aardvark is running in the zoo\", \"id\": 4},\r\n",
      "{\"content\": \"Earlier today, scientists from Cambodia discovered\", \"id\": 5},\r\n",
      "{\"content\": \"Earlier today, scientists from Bangladesh discovered\", \"id\": 6},\r\n",
      "{\"content\": \"Earlier today, scientists from Spain discovered\", \"id\": 7},\r\n",
      "{\"content\": \"Earlier today, scientists from Solomon Islands discovered\", \"id\": 8},\r\n",
      "{\"content\": \"Earlier today, scientists from Eritrea discovered\", \"id\": 9},\r\n",
      "{\"content\": \"Earlier today, scientists from Togo discovered\", \"id\": 10},\r\n",
      "{\"content\": \"Earlier today, scientists from Mexico discovered\", \"id\": 11},\r\n",
      "{\"content\": \"Earlier today, scientists from Sudan discovered\", \"id\": 12},\r\n",
      "{\"content\": \"Earlier today, scientists from Saint Kitts and Nevis discovered\", \"id\": 13},\r\n",
      "{\"content\": \"Earlier today, scientists from India discovered\", \"id\": 14},\r\n",
      "{\"content\": \"Carolyn is my neighbor.\", \"id\": 15},\r\n",
      "{\"content\": \"John is my neighbor.\", \"id\": 16},\r\n",
      "{\"content\": \"George is my neighbor.\", \"id\": 17},\r\n",
      "{\"content\": \"Lawrence is my neighbor.\", \"id\": 18},\r\n",
      "{\"content\": \"Linda is my neighbor.\", \"id\": 19},\r\n",
      "{\"content\": \"Steve is my neighbor.\", \"id\": 20},\r\n",
      "{\"content\": \"Judy is my neighbor.\", \"id\": 21},\r\n",
      "{\"content\": \"Alfred is my neighbor.\", \"id\": 22},\r\n",
      "{\"content\": \"Patricia is my neighbor.\", \"id\": 23},\r\n",
      "{\"content\": \"Jill is my neighbor.\", \"id\": 24}]}"
     ]
    }
   ],
   "source": [
    "cat 'same_token_suite.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the test suite JSON\n",
    "The JSON file we created can be imported back into a Python object by using `json.load()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'content': 'The dog is running in the zoo', 'id': 1},\n",
       " {'content': 'The cat is running in the zoo', 'id': 2},\n",
       " {'content': 'The giraffe is running in the zoo', 'id': 3}]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "f = open('same_token_suite.json', 'r')\n",
    "suite_dict = json.load(f)\n",
    "f.close()\n",
    "suite_dict['examples'][0:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating predictions from the loaded data\n",
    "Our data has been loaded into a variable named `suite_dict`. Now we can read each example from `suite_dict` and generate the predictions. Each prediction will be written to another file named `same_token_suite_predictions.json`, which will be sent to Checklist to evaluate the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('same_token_suite_predictions.json', 'w') as f:\n",
    "    for example in suite_dict['examples']:\n",
    "        prediction = generate_sentence(example['content'])\n",
    "        prediction = prediction.replace('\"', '\\\"')\n",
    "        f.write(json.dumps({'prediction': prediction, 'id': example['id']}) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"prediction\": \". Just take care of a dead dog!\\n\\nWe'll do our best to make this possible. We are sure this will be a one-time event (not the main event) that will last no more\", \"id\": 1}\r\n",
      "{\"prediction\": \"'s water purifier with his family \\u2013 from where they plan to raise goats \\u2013 and he'll probably get better at sleeping next to them, he told me. The kitten, called Bambu is only about eight\", \"id\": 2}\r\n",
      "{\"prediction\": \". It's the most endangered vertebrate in North America, and has a population of only 12,000. And the animal's diet is similar to that of all the other mammals on the planet - meat:\", \"id\": 3}\r\n",
      "{\"prediction\": \" with the other monkeys and all of the animals, and we want to make sure all animals are well cared for, safe and happy.\\n\\nAs soon as we finish, we'll send it out\", \"id\": 4}\r\n",
      "{\"prediction\": \" an ancient fossil in the Khao Chinchon Forest that was recovered in 1994. The remains are thought to date back to 868 B.C., a period of heavy agriculture and trade through which the Haus\", \"id\": 5}\r\n",
      "{\"prediction\": \" that certain genes have a profound effect on the immune system, resulting in an increased rate of disease.\\n\\nThe study, published in the journal Nature, revealed that DNA in bacteria are able to break down and bind\", \"id\": 6}\r\n",
      "{\"prediction\": \" that a key part of the body's immune system produces antibodies that act as \\\"anti-antibodies\\\"\\u2014which are required to keep out pathogens that harm cells. The two antibodies, which the researchers say are in\", \"id\": 7}\r\n",
      "{\"prediction\": \" the first ever confirmed species of human sperm by testing them in a new laboratory. This discovery could change life, says one of the researchers, who also has a Ph.D. in comparative anatomy.\\n\\n\", \"id\": 8}\r\n",
      "{\"prediction\": \" that some of Eritrean colonists' DNA matched their own for the first time.\\n\\n\\nThe DNA in a collection containing seven human-pig remains found on the coast of Sudan is almost identical to a\", \"id\": 9}\r\n",
      "{\"prediction\": \" that they had detected a small fragment of the Earth's mantle at exactly the right volume, more than seven times the amount found in the previously known mass fraction and well below the original amount of 9.8 times\", \"id\": 10}\r\n",
      "{\"prediction\": \" a new chemical they have found responsible for the body of E.coli that was used to control the immune system in a lab at Arizona State University.\\n\\n\\\"Today, when it works we see immune molecules\", \"id\": 11}\r\n",
      "{\"prediction\": \" that the parasite was the first organism to use its unique host organism, Gyeonphaenoicollis, as its own colony.\\n\\nThat's right -- both individuals from the same species, the genus\", \"id\": 12}\r\n",
      "{\"prediction\": \" that they could observe a comet's orbit through an electromagnetic field.\\n\\nTheir measurements show comet comets must have orbits around the Sun, a mechanism that could allow scientists to detect the faint\", \"id\": 13}\r\n",
      "{\"prediction\": \" the first ever structure that appears to correspond to the origins of DNA from the remains of a man, according to a first-of-its-kind study in the Proceedings of the National Academy of Sciences (PNAS\", \"id\": 14}\r\n",
      "{\"prediction\": \" I am her neighbor and she loves me and is in love with me. She just wants to make sure everybody in this community understands I don't have any assets that can be used to take care of her business, to\", \"id\": 15}\r\n",
      "{\"prediction\": \" I'm still here. He had to pull the blanket over me. It's no good.\\n\\nI'm at this. Please tell me some things. One thing, thank you. Did his sister steal your car last\", \"id\": 16}\r\n",
      "{\"prediction\": \" My grandmother moved in with me when I was a kid and my mom has been going a fair bit of shopping since she can, so she has taken care of me from the moment I left my parents' home to the point\", \"id\": 17}\r\n",
      "{\"prediction\": \" He is good, and he takes up to 18 years to be back, which is a great honor and a pleasure to sit with, because I am just glad I live. I will see him later tonight or tonight and\", \"id\": 18}\r\n",
      "{\"prediction\": \" My sister and I both have kids and we are married. She's already got two beautiful girls.\\n\\nMy brother and he have been together. Now, as soon as he had that baby, we started walking around\", \"id\": 19}\r\n",
      "{\"prediction\": \" I can't let you down, but you need to see my face.\\n\\nI'll have some fun with you at the bar later. Cheers.\", \"id\": 20}\r\n",
      "{\"prediction\": \" So I am with my friends and family. In this case, he has a problem. His friend gets mad. One of his friends can't walk with his feet because she has Parkinson's. And since my friend has\", \"id\": 21}\r\n",
      "{\"prediction\": \" He's just as smart as I am, has a fair sized budget, and is a fan of the Marvel Universe. And he just loves to sing.\\n\\nBut if the question becomes whether or not he has\", \"id\": 22}\r\n",
      "{\"prediction\": \" She loves you. I'm an immigrant but also a hard-working mom. Now my little brother is 9. He's very smart, but he doesn't understand anything. So for 10 years that's our life.\", \"id\": 23}\r\n",
      "{\"prediction\": \" We come home after lunch every night to smoke the sweet stuff. I always thought that smokers needed cigarettes.\\n\\nWe smoke a lot in Texas, but don't know about the rest of the state. It's really\", \"id\": 24}\r\n"
     ]
    }
   ],
   "source": [
    "cat 'same_token_suite_predictions.json'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading test results from file\n",
    "TestSuite has a `run_from_file()` function that reads the predictions line by line from a file. The `format_fn` parameter is used to parse each line of the file. Our format function, `read_json_prediction()`, converts the JSON object into a tuple of `(predicted text, confidence)`. We don't care about confidence values here, so we will just set confidence to 1 for every prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_json_prediction(x):\n",
    "    test_output = json.loads(x)\n",
    "    return (test_output['prediction'], 1)\n",
    "suite.run_from_file('same_token_suite_predictions.json', format_fn = read_json_prediction, overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Same Token Prediction\n",
      "\n",
      "Same animal in response\n",
      "Test cases:      4\n",
      "Fails (rate):    3 (75.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      The cat is running in the zoo\n",
      "Completion:      's water purifier with his family – from where they plan to raise goats – and he'll probably get better at sleeping next to them, he told me. The kitten, called Bambu is only about eight\n",
      "----\n",
      "Prompt:      The aardvark is running in the zoo\n",
      "Completion:       with the other monkeys and all of the animals, and we want to make sure all animals are well cared for, safe and happy.\n",
      "\n",
      "As soon as we finish, we'll send it out\n",
      "----\n",
      "Prompt:      The giraffe is running in the zoo\n",
      "Completion:      . It's the most endangered vertebrate in North America, and has a population of only 12,000. And the animal's diet is similar to that of all the other mammals on the planet - meat:\n",
      "----\n",
      "\n",
      "\n",
      "Same country in response\n",
      "Test cases:      10\n",
      "Fails (rate):    9 (90.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Earlier today, scientists from Spain discovered\n",
      "Completion:       that a key part of the body's immune system produces antibodies that act as \"anti-antibodies\"—which are required to keep out pathogens that harm cells. The two antibodies, which the researchers say are in\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Sudan discovered\n",
      "Completion:       that the parasite was the first organism to use its unique host organism, Gyeonphaenoicollis, as its own colony.\n",
      "\n",
      "That's right -- both individuals from the same species, the genus\n",
      "----\n",
      "Prompt:      Earlier today, scientists from Saint Kitts and Nevis discovered\n",
      "Completion:       that they could observe a comet's orbit through an electromagnetic field.\n",
      "\n",
      "Their measurements show comet comets must have orbits around the Sun, a mechanism that could allow scientists to detect the faint\n",
      "----\n",
      "\n",
      "\n",
      "Same person in response\n",
      "Test cases:      10\n",
      "Fails (rate):    10 (100.0%)\n",
      "\n",
      "Example fails:\n",
      "Prompt:      Steve is my neighbor.\n",
      "Completion:       I can't let you down, but you need to see my face.\n",
      "\n",
      "I'll have some fun with you at the bar later. Cheers.\n",
      "----\n",
      "Prompt:      Alfred is my neighbor.\n",
      "Completion:       He's just as smart as I am, has a fair sized budget, and is a fan of the Marvel Universe. And he just loves to sing.\n",
      "\n",
      "But if the question becomes whether or not he has\n",
      "----\n",
      "Prompt:      Jill is my neighbor.\n",
      "Completion:       We come home after lunch every night to smoke the sweet stuff. I always thought that smokers needed cigarettes.\n",
      "\n",
      "We smoke a lot in Texas, but don't know about the rest of the state. It's really\n",
      "----\n",
      "\n",
      "\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "suite.summary(format_example_fn = format_example)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Please wait as we prepare the table data...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "091d107856424cfe87f91d2fea0d0d4a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "SuiteSummarizer(stats={'npassed': 0, 'nfailed': 0, 'nfiltered': 0}, test_infos=[{'name': 'Same animal in respo…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "suite.visual_summary_table()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
